# batch_inference_4class.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
VEHICLE_MODEL_PATH = r"D:\CodeCNN\intel\vehicle-detection-0202\FP32\vehicle-detection-0202.xml"
ROAD_MODEL_PATH = r"D:\CodeCNN\intel\road-segmentation-adas-0001\FP32\road-segmentation-adas-0001.xml"
INPUT_FOLDER = r"H:\xiaomi\test"  # 输入图片文件夹
OUTPUT_FOLDER = r"detect"  # 输出结果文件夹

# 创建输出目录
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# -----------------------------
# 初始化 OpenVINO
# -----------------------------
core = Core()
device = "CPU"  # 可选 "GPU"（若支持）

# -----------------------------
# 1. 加载车辆检测模型
# -----------------------------
print("Loading vehicle detection model...")
vehicle_model = core.read_model(model=VEHICLE_MODEL_PATH)
vehicle_compiled = core.compile_model(model=vehicle_model, device_name=device)
vehicle_input_layer = vehicle_compiled.input(0)
vehicle_output_layer = vehicle_compiled.output(0)

# 获取输入尺寸
_, _, h_det, w_det = vehicle_input_layer.shape
print(f"Vehicle detection input shape: {h_det}x{w_det}")

# -----------------------------
# 2. 加载道路分割模型（4类）
# -----------------------------
print("Loading 4-class road segmentation model...")
road_model = core.read_model(model=ROAD_MODEL_PATH)
road_compiled = core.compile_model(model=road_model, device_name=device)
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

# 输出形状应为 [1, 4, H, W]
_, num_classes, h_seg, w_seg = road_output_layer.shape
print(f"Road segmentation output classes: {num_classes} (should be 4)")
if num_classes != 4:
    print("⚠️ 模型输出不是4类！请确认是否使用了多分类版本（如 -1024x768）")
    exit(1)

print(f"Road segmentation input shape: {h_seg}x{w_seg}")

# -----------------------------
# 3. 获取 A 文件夹中所有图片
# -----------------------------
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
image_files = []
for root, dirs, files in os.walk(INPUT_FOLDER):
    for file in files:
        if os.path.splitext(file)[1].lower() in image_extensions:
            image_files.append(os.path.join(root, file))

if not image_files:
    print(f"❌ 没有在 '{INPUT_FOLDER}' 中找到图片！")
    exit(1)

print(f"✅ 发现 {len(image_files)} 张图片，开始批量处理...")

# -----------------------------
# 4. 定义 4 类颜色映射（BGR 格式）- 背景透明
# -----------------------------
# 0: background → 透明（不绘制）
# 1: road → 绿色
# 2: curb → 蓝色（路缘）
# 3: road markings → 红色（车道线）
COLOR_MAP = {
    0: None,  # 背景 - 透明
    1: (0, 255, 0),  # green - 道路
    2: (255, 0, 0),  # blue - 路缘
    3: (0, 0, 255),  # red - 道路标记
}


# -----------------------------
# 5. 优化版车辆压线检测函数（方案1：中心区域检测）
# -----------------------------
def check_lane_violation_optimized(vehicle_bbox, lane_marking_mask, confidence_threshold=0.01, side_margin_ratio=0.2):
    """
    优化版压线检测 - 只检测车辆底部中心区域，避开两侧

    Args:
        vehicle_bbox: 车辆边界框 [x_min, y_min, x_max, y_max]
        lane_marking_mask: 道路标记掩码（二值图像）
        confidence_threshold: 重叠区域阈值比例
        side_margin_ratio: 两侧排除比例（0.2表示排除左右各20%）

    Returns:
        is_violation: 是否压线
        overlap_ratio: 重叠区域比例
        overlap_pixels: 重叠像素数
        vehicle_bottom_bbox: 检测的底部中心区域坐标
        overlap_mask: 重叠区域掩码
    """
    x_min, y_min, x_max, y_max = vehicle_bbox

    # 提取车辆底部区域（重点关注车轮接触区域）
    bottom_height_ratio = 0.15  # 检测底部15%区域
    bottom_y_min = int(y_max - (y_max - y_min) * bottom_height_ratio)

    # 排除两侧区域，只检测中心部分（避免包含旁边的实线）
    side_margin = int((x_max - x_min) * side_margin_ratio)
    center_x_min = x_min + side_margin
    center_x_max = x_max - side_margin

    # 确保坐标在图像范围内
    bottom_y_min = max(bottom_y_min, 0)
    bottom_y_max = min(y_max, lane_marking_mask.shape[0])
    center_x_min = max(center_x_min, 0)
    center_x_max = min(center_x_max, lane_marking_mask.shape[1])

    # 如果中心区域太小，使用完整宽度
    if center_x_max <= center_x_min or (center_x_max - center_x_min) < 10:
        center_x_min = x_min
        center_x_max = x_max
        print("    ⚠️ 中心区域过小，使用完整宽度检测")

    vehicle_bottom_bbox = [center_x_min, bottom_y_min, center_x_max, bottom_y_max]

    # 创建车辆底部中心区域的掩码
    vehicle_bottom_mask = np.zeros_like(lane_marking_mask)
    if bottom_y_max > bottom_y_min and center_x_max > center_x_min:
        vehicle_bottom_mask[bottom_y_min:bottom_y_max, center_x_min:center_x_max] = 1

    # 计算重叠区域
    overlap_mask = vehicle_bottom_mask & lane_marking_mask
    overlap_pixels = np.sum(overlap_mask)

    # 计算车辆底部中心区域的总像素数
    vehicle_bottom_area = (center_x_max - center_x_min) * (bottom_y_max - bottom_y_min)

    if vehicle_bottom_area > 0:
        overlap_ratio = overlap_pixels / vehicle_bottom_area
    else:
        overlap_ratio = 0

    # 详细的调试信息
    print(f"    车辆位置: [{x_min}, {y_min}, {x_max}, {y_max}]")
    print(f"    检测区域: 底部{int(bottom_height_ratio * 100)}% 中心{int((1 - 2 * side_margin_ratio) * 100)}%")
    print(f"    中心区域: [{center_x_min}, {bottom_y_min}, {center_x_max}, {bottom_y_max}]")
    print(f"    区域面积: {vehicle_bottom_area}px, 重叠像素: {overlap_pixels}px")
    print(f"    重叠比例: {overlap_ratio:.3f} ({overlap_ratio * 100:.1f}%)")
    print(f"    检测阈值: {confidence_threshold} ({confidence_threshold * 100:.1f}%)")

    # 判断是否压线
    is_violation = overlap_ratio > confidence_threshold

    if is_violation:
        print(f"    🚨 检测到压线违规!")
    else:
        if overlap_pixels > 0:
            print(f"    ✅ 有重叠但未达阈值")
        else:
            print(f"    ✅ 无重叠")

    return is_violation, overlap_ratio, overlap_pixels, vehicle_bottom_bbox, overlap_mask


# -----------------------------
# 6. 批量处理每张图片
# -----------------------------
violation_stats = []  # 存储违规统计信息

for i, img_path in enumerate(image_files, 1):
    print(f"\n--- 处理第 {i}/{len(image_files)} 张: {os.path.basename(img_path)} ---")

    # 读取图像
    image = cv2.imread(img_path)
    if image is None:
        print(f"⚠️ 无法加载图像: {img_path}")
        continue

    orig_h, orig_w = image.shape[:2]
    print(f"Original size: {orig_w}x{orig_h}")

    # -----------------------------
    # 7. 车辆检测推理
    # -----------------------------
    resized_det = cv2.resize(image, (w_det, h_det))
    input_det = np.expand_dims(resized_det.transpose(2, 0, 1), axis=0)  # HWC -> NCHW
    detections = vehicle_compiled([input_det])[vehicle_output_layer]

    # -----------------------------
    # 8. 道路分割推理（4类）
    # -----------------------------
    resized_seg = cv2.resize(image, (w_seg, h_seg))
    input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)

    # 推理
    infer_result = road_compiled([input_seg])
    tensor_output = infer_result[road_output_layer]

    # 转为 NumPy 并处理
    seg_probs = np.array(tensor_output)  # shape: [1, 4, H, W]
    seg_class = np.argmax(seg_probs[0], axis=0).astype(np.uint8)
    seg_class = np.ascontiguousarray(seg_class)  # 确保内存连续

    # 调试信息
    print(f"seg_class dtype: {seg_class.dtype}, shape: {seg_class.shape}")
    unique_vals, counts = np.unique(seg_class, return_counts=True)
    print(f"分割结果类别分布:")
    for val, count in zip(unique_vals, counts):
        print(f"  类别 {val}: {count} 像素 ({count / seg_class.size * 100:.2f}%)")

    # -----------------------------
    # 9. 将分割结果缩放到原图尺寸
    # -----------------------------
    seg_class_resized = cv2.resize(seg_class, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)

    # -----------------------------
    # 10. 生成彩色语义分割掩码
    # -----------------------------
    seg_color_mask = np.zeros_like(image)
    non_bg_mask = np.zeros((orig_h, orig_w), dtype=bool)  # 非背景区域掩码

    for cls, color in COLOR_MAP.items():
        if color is None:  # 跳过背景
            continue
        mask = (seg_class_resized == cls)
        mask_pixels = np.sum(mask)
        print(f"类别 {cls} 的像素数: {mask_pixels}")

        if mask_pixels > 0:
            seg_color_mask[mask] = color
            non_bg_mask = non_bg_mask | mask  # 累积非背景区域
            print(f"  → 为类别 {cls} 应用颜色 {color}")

    # 检查背景处理
    bg_mask = (seg_class_resized == 0)
    bg_pixels = np.sum(bg_mask)
    print(f"背景像素数: {bg_pixels}")
    print(f"非背景区域像素数: {np.sum(non_bg_mask)}")

    # -----------------------------
    # 11. 创建道路标记掩码（用于压线检测）
    # -----------------------------
    lane_marking_mask = (seg_class_resized == 3)  # 类别3是道路标记
    lane_marking_pixels = np.sum(lane_marking_mask)
    print(f"道路标记像素数: {lane_marking_pixels}")

    # -----------------------------
    # 12. 车辆压线检测和绘制
    # -----------------------------
    output_image = image.copy()
    car_count = 0
    violation_count = 0

    # 创建检测区域可视化图层
    detection_overlay = output_image.copy()

    # 首先绘制所有检测区域（黄色框和半透明背景）
    for det_idx, det in enumerate(detections[0, 0]):
        conf = float(det[2])
        if conf > 0.5:
            x_min = int(det[3] * orig_w)
            y_min = int(det[4] * orig_h)
            x_max = int(det[5] * orig_w)
            y_max = int(det[6] * orig_h)

            vehicle_bbox = [x_min, y_min, x_max, y_max]

            # 使用优化版压线检测函数
            is_violation, overlap_ratio, overlap_pixels, bottom_bbox, overlap_mask = check_lane_violation_optimized(
                vehicle_bbox, lane_marking_mask, confidence_threshold=0.001, side_margin_ratio=0.2
            )

            # 绘制检测底部中心区域的半透明背景
            bx_min, by_min, bx_max, by_max = bottom_bbox

            # 根据是否违规选择颜色
            if is_violation:
                overlay_color = (0, 0, 255)  # 红色 - 违规
            else:
                overlay_color = (0, 255, 255)  # 黄色 - 正常

            alpha = 0.3  # 透明度

            # 绘制半透明矩形（只在有效区域内）
            if by_max > by_min and bx_max > bx_min:
                overlay_area = detection_overlay[by_min:by_max, bx_min:bx_max]
                if overlay_area.size > 0:
                    color_overlay = np.full_like(overlay_area, overlay_color)
                    cv2.addWeighted(color_overlay, alpha, overlay_area, 1 - alpha, 0, overlay_area)

            # 绘制边框
            border_color = (0, 0, 255) if is_violation else (0, 255, 255)
            cv2.rectangle(detection_overlay, (bx_min, by_min), (bx_max, by_max), border_color, 2)

            # 添加检测区域标签
            label_bg_color = border_color
            label_text_color = (0, 0, 0)  # 白色文字

            text1 = f"Center Area {det_idx + 1}"
            text2 = f"Overlap: {overlap_pixels}px"
            text3 = f"Ratio: {overlap_ratio:.1%}"

            (text1_width, text1_height), _ = cv2.getTextSize(text1, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
            (text2_width, text2_height), _ = cv2.getTextSize(text2, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
            (text3_width, text3_height), _ = cv2.getTextSize(text3, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)

            text_width = max(text1_width, text2_width, text3_width)
            text_height = text1_height + text2_height + text3_height + 10

            # 绘制文本背景
            text_bg_x1 = bx_min
            text_bg_y1 = by_min - text_height - 5
            text_bg_x2 = bx_min + text_width + 6
            text_bg_y2 = by_min - 5

            # 确保文本背景在图像范围内
            if text_bg_y1 < 0:
                text_bg_y1 = by_max + 5
                text_bg_y2 = text_bg_y1 + text_height + 5

            cv2.rectangle(detection_overlay, (text_bg_x1, text_bg_y1),
                          (text_bg_x2, text_bg_y2), label_bg_color, -1)
            cv2.rectangle(detection_overlay, (text_bg_x1, text_bg_y1),
                          (text_bg_x2, text_bg_y2), (0, 0, 0), 1)

            # 绘制文本
            cv2.putText(detection_overlay, text1, (bx_min + 3, text_bg_y1 + text1_height + 3),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, label_text_color, 1)
            cv2.putText(detection_overlay, text2, (bx_min + 3, text_bg_y1 + text1_height + text2_height + 6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, label_text_color, 1)
            cv2.putText(detection_overlay, text3, (bx_min + 3, text_bg_y1 + text_height - 3),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, label_text_color, 1)

    # 将检测区域图层叠加到输出图像上
    output_image = detection_overlay

    # 然后绘制车辆边界框和标签
    for det_idx, det in enumerate(detections[0, 0]):
        conf = float(det[2])
        if conf > 0.5:
            x_min = int(det[3] * orig_w)
            y_min = int(det[4] * orig_h)
            x_max = int(det[5] * orig_w)
            y_max = int(det[6] * orig_h)

            vehicle_bbox = [x_min, y_min, x_max, y_max]

            # 使用优化版压线检测
            is_violation, overlap_ratio, overlap_pixels, bottom_bbox, overlap_mask = check_lane_violation_optimized(
                vehicle_bbox, lane_marking_mask, confidence_threshold=0.001, side_margin_ratio=0.2
            )

            # 根据检测结果设置颜色和标签
            if is_violation:
                bbox_color = (0, 0, 255)  # 红色 - 违规
                label = f"VIOLATION {det_idx + 1}: {overlap_ratio:.1%}"
                violation_count += 1
                print(f"🚨 车辆 {det_idx + 1} 压线! 重叠比例: {overlap_ratio:.3f}")
            else:
                bbox_color = (0, 255, 0)  # 绿色 - 正常
                label = f"Car {det_idx + 1}: {conf:.2f}"

            # 绘制车辆边界框（较粗的线）
            cv2.rectangle(output_image, (x_min, y_min), (x_max, y_max), bbox_color, 3)

            # 绘制车辆标签背景
            label_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
            label_bg_x1 = x_min
            label_bg_y1 = y_min - label_size[1] - 10
            label_bg_x2 = x_min + label_size[0] + 10
            label_bg_y2 = y_min

            # 确保标签在图像范围内
            if label_bg_y1 < 0:
                label_bg_y1 = y_max
                label_bg_y2 = y_max + label_size[1] + 10
                label_pos = (x_min + 5, y_max + label_size[1] + 5)
            else:
                label_pos = (x_min + 5, y_min - 5)

            cv2.rectangle(output_image,
                          (label_bg_x1, label_bg_y1),
                          (label_bg_x2, label_bg_y2),
                          bbox_color, -1)

            # 绘制车辆标签文字
            cv2.putText(output_image, label, label_pos,
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

            car_count += 1

    print(f"检测到车辆数: {car_count}, 压线车辆数: {violation_count}")

    # 记录统计信息
    violation_stats.append({
        'filename': os.path.basename(img_path),
        'total_vehicles': car_count,
        'violations': violation_count,
        'violation_ratio': violation_count / max(car_count, 1)
    })

    # -----------------------------
    # 13. 叠加语义分割（只在非背景区域叠加）
    # -----------------------------
    overlay = output_image.copy()

    # 只在非背景区域进行叠加
    if np.any(non_bg_mask):
        overlay[non_bg_mask] = cv2.addWeighted(
            output_image[non_bg_mask], 0.6,
            seg_color_mask[non_bg_mask], 0.4, 0
        )
        print("已应用语义分割叠加")
    else:
        print("没有检测到道路相关区域，保持原图")

    # 在图像上添加统计信息
    stats_text = f"Vehicles: {car_count}  Violations: {violation_count}"
    cv2.putText(overlay, stats_text, (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 3)
    cv2.putText(overlay, stats_text, (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1)

    # -----------------------------
    # 14. 保存结果到 detect 目录
    # -----------------------------
    base_name = os.path.splitext(os.path.basename(img_path))[0]
    output_filename = f"{base_name}_result.jpg"
    output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    # 如果文件已存在，添加时间戳避免覆盖
    if os.path.exists(output_path):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        output_filename = f"{base_name}_result_{timestamp}.jpg"
        output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    cv2.imwrite(output_path, overlay)
    print(f"✅ 结果已保存至: {output_path}")

    # -----------------------------
    # 15. （可选）显示最后一张图
    # -----------------------------
    if i == len(image_files):  # 只显示最后一张
        def resize_to_fit(img, target_w=1360, target_h=768):
            h, w = img.shape[:2]
            scale = min(target_w / w, target_h / h)
            new_w, new_h = int(w * scale), int(h * scale)
            resized = cv2.resize(img, (new_w, new_h))
            canvas = np.zeros((target_h, target_w, 3), dtype=np.uint8)
            y_pad = (target_h - new_h) // 2
            x_pad = (target_w - new_w) // 2
            canvas[y_pad:y_pad + new_h, x_pad:x_pad + new_w] = resized
            return canvas


        display_img = resize_to_fit(overlay)
        cv2.namedWindow("Batch Result", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Batch Result", 1360, 768)
        cv2.imshow("Batch Result", display_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

# -----------------------------
# 16. 输出压线统计报告
# -----------------------------
print(f"\n{'=' * 50}")
print("🚗 车辆压线检测统计报告（优化版）")
print(f"{'=' * 50}")

total_vehicles = sum(stat['total_vehicles'] for stat in violation_stats)
total_violations = sum(stat['violations'] for stat in violation_stats)

print(f"总处理图片数: {len(violation_stats)}")
print(f"总检测车辆数: {total_vehicles}")
print(f"总压线违规数: {total_violations}")
print(f"总体违规比例: {total_violations / max(total_vehicles, 1) * 100:.1f}%")

print(f"\n详细统计:")
for stat in violation_stats:
    if stat['total_vehicles'] > 0:
        print(
            f"  {stat['filename']}: {stat['total_vehicles']}车, {stat['violations']}违规 ({stat['violation_ratio'] * 100:.1f}%)")

print(f"\n🎉 所有 {len(image_files)} 张图片处理完成！结果保存在 '{OUTPUT_FOLDER}' 目录中。")